"""
This module contains FunctionCall and Command classes.  These classes
support the following methods:
    run() -> (return_code, return_value)
        Invoke the function/shell command and return the command-style return
        code and and optional extra return data.
    to_dict() -> dict
        Return a dict representation of the object that can be marshalled
        with json.

This module has a from_dictionary function that will convert from
a dict representation to a corresponding FunctionCall or Command object.

When creating a Command object you can specify a dictionary of environment
variables that apply only to that command.  However, this module also has
functions to set environment variables that apply to subsequent commands
that are created.

Example 1:
    # "VAR1" will be set when both command1 and command2 run.
    command.set_environ_var("VAR1", "1")
    command1 = command.Command(...)
    command2 = command.Command(...)

Example 2:
    # "VAR1" will be set whe command1 and command2 run but not command3.
    with command.push_environ() as environ_guard:
        command.set_environ_var("VAR1", "1")
        command1 = command.Command(...)
        command2 = command.Command(...)
    command3 = command.Command(...)
"""

import sys
import os
import types
import subprocess
import platform
import shlex
import traceback
import marshal
import re
import base64
import ctypes

if platform.system() == "Windows":
    import msvcrt
    import ctypes
    import ctypes.wintypes

    win_PeekNamedPipe = ctypes.WINFUNCTYPE(
            ctypes.wintypes.BOOL, # return value
            ctypes.wintypes.HANDLE, # named pipe
            ctypes.wintypes.LPVOID, # buffer for output
            ctypes.wintypes.DWORD, # output buffer size
            ctypes.POINTER(ctypes.wintypes.DWORD), # bytes read
            ctypes.POINTER(ctypes.wintypes.DWORD), # total bytes available
            ctypes.POINTER(ctypes.wintypes.DWORD), # bytes left in this message
        )(("PeekNamedPipe", ctypes.windll.kernel32))

else:
    import select

#-----------------------------------------------------------------------------
# Code to unmarshal from a dictionary:

def from_dict(dictionary):
    """Given a dictionary representing a marshalled FunctionCall or Command
    object, return an instance of the corresponding object.
    """
    return (FunctionCall.unmarshal(dictionary["marshalled_function_call"])
        if "marshalled_function_call" in dictionary
        else Command.unmarshal(dictionary))

#-----------------------------------------------------------------------------
# Code to set environment variables when commands run:

_environ_stack = [{}]

def set_environ_var(name, value=""):
    """Schedule a variable to be set in the environment whenever a command
    runs.  Whenever a Command object is created it will use the current set
    of environment variables.
    """
    _environ_stack[-1][name] = value

def push_environ():
    """Push a new set of environment variable settings.  You can then call
    set_environ_var a number of times, create Command objects, and undo the
    environment variables by calling pop_environ.

    Optionally, you can use the return value in a with statement to have the
    environment variable set automatically popped.
        with command.push_environ() as environ_guard:
            command.set_environ_var("FOO", "1")
            command.Command(...)
    """
    _environ_stack.append({})
    return _PushedEnvironGuard()

def pop_environ():
    """Pop a set of environment variables that was earlier pushed."""
    _environ_stack.pop()

class _PushedEnvironGuard(object):
    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        pop_environ()

def _collapsed_environ():
    """Return the combined dictionary of all environment variables in all the
    pushed environment sets.
    """
    result_environ = {}
    for environ in _environ_stack:
        result_environ.update(environ)
    return result_environ

#-----------------------------------------------------------------------------
# Code for function calls:

_create_cell = ctypes.pythonapi.PyCell_New
_create_cell.restype = ctypes.py_object
_create_cell.argtypes = (ctypes.py_object,)

class FunctionCall(object):
    """This class allows you to pass function calls between separate Python
    processes.

    Below is an example of how to use this module.  From one process:
        s = command.FunctionCall(
            "rm /file/to/remove", "__import__('os').unlink",
            "/file/to/remove").marshal()
        ...send s to another process...
    From the other process:
        s = ...receive marshalled string from the first process...
        function_call = command.FunctionCall.unmarshal(s)
        print function_call
        function_call()
    """
    @classmethod
    def unmarshal(cls, string, globals=None):
        if globals is None:
            globals = {"__builtins__": __builtins__}

        description, code, name, default_args, closure_values, args, kwargs = (
            marshal.loads(base64.b64decode(string)))
        closure = (
            tuple(_create_cell(closure_value)
                for closure_value in closure_values)
            if closure_values is not None else None)
        function = types.FunctionType(
            code, globals, name, default_args, closure)
        return cls(description, function, *args, **kwargs)

    def __init__(self, description, function, *args, **kwargs):
        self.description = description
        self.function = _marshallable_function(function)
        self.args = args
        self.kwargs = kwargs

    @property
    def command(self):
        """Return a command line corresponding to this function call."""
        return self.description

    def __repr__(self):
        return "<FunctionCall %s>" % self.description

    def to_dict(self):
        return {"marshalled_function_call": self.marshal()}

    def run(self):
        # TODO: Add a callback parameter that is invoked as output becomes
        #       available, so we can have different front-ends to display the
        #       output.
        """Return a (return_code, return_value) tuple.  return_code is
        a program's exit code number (0 for success).  return_value provides
        a way to return extra information, and can be None to indicate there
        is no extra information.
        """
        try:
            return_value = self.function(*self.args, **self.kwargs)
        except:
            traceback.print_exc()
            return 1, None

        return 0, return_value

    def marshal(self):
        closure_values = (
            tuple(cell.cell_contents for cell in self.function.func_closure)
            if self.function.func_closure is not None
            else None)
        return base64.b64encode(marshal.dumps((
            self.description, self.function.func_code, self.function.__name__,
            self.function.func_defaults, closure_values,
            self.args, self.kwargs)))

def _marshallable_function(function):
    """If "function" is a function written in Python, just return it.  If it is
    a string containing the name of a C extension function, return a function
    that wraps it.  Otherwise, if it a C extension function, raise ValueError.
    """
    if not isinstance(function, str):
        if hasattr(function, "func_code"):
            return function
        raise ValueError("You cannot pass in a C extension function."
            "  Instead, pass in the name of that function.")

    namespace = {}
    eval(compile(
        "def wrapped_function(*args, **kwargs):\n"
        "    return %s(*args, **kwargs)\n" % function,
        "<command.py (wrapper)>", "exec"), namespace)
    return namespace["wrapped_function"]

#-----------------------------------------------------------------------------
# Code to run command-line programs:

class Command(object):
    @classmethod
    def unmarshal(cls, dictionary):
        dictionary = dictionary.copy()
        for filter_name in "stdout_filter", "stderr_filter":
            dictionary[filter_name] = cls._unmarshal_filter(
                dictionary[filter_name])

        return cls(**dictionary)

    @classmethod
    def _unmarshal_filter(cls, marshalled_data):
        return (None if marshalled_data is None
            else FunctionCall.unmarshal(marshalled_data).function)

    def __init__(
            self, command, description=None, use_shell=False,
            stdout_filter=None, stderr_filter=None, environ=None):
        """Represents a command line that can be run.

        command:
            Either a string representing the command-line or an argv list.

        description:
            A description of the command.  If None, the command line will be
            used as the description.

        use_shell: 
            Runs the command by spawning a bash shell if set to True.

        stdout_filter
        stderr_filter
            If not None, these filters are callable objects that take a line
            of output (including the newline) and return a line of filtered
            output.  By returning an empty string, they will remove that line
            from the output.  The filter is called with an empty string to
            signal that the end of the output has been reached.
        """
        # If command is a sequence of command-line arguments instead of a
        # string containing a single command line, form a single command line
        # from the arguments.
        if not isinstance(command, str):
            command = form_command_line(command)

        self.command = command
        self._description = description
        self.use_shell = use_shell
        self.stdout_filter = stdout_filter
        self.stderr_filter = stderr_filter

        # Remember the current environment variable settings so we can set them
        # when the command runs.  Note that when we're marshalled to the client
        # process and constructed there won't be any environment variable
        # settings, but the full set of environment variables will be passed to
        # the constructor.
        self.environ = _collapsed_environ()
        if environ is not None:
            self.environ.update(environ)

    def __repr__(self):
        return "<Command %s>" % self.command

    @property
    def description(self):
        return (self._description if self._description is not None
            else self.command)

    def to_dict(self):
        return {
            "command": self.command,
            "description": self.description,
            "use_shell": self.use_shell,
            "stdout_filter": self._marshal_filter(self.stdout_filter),
            "stderr_filter": self._marshal_filter(self.stderr_filter),
            "environ": self.environ,
        }

    def _marshal_filter(self, output_filter):
        return (None if output_filter is None
            else FunctionCall("output_filter", output_filter).marshal())

    def run(self):
        # TODO: Add an example of using the filters as a way to have a
        #       different front end that displays the output.
        """Return a (return_code, return_value) tuple.  return_code is
        a program's exit code number (0 for success).  return_value provides
        a way to return extra information, and can be None to indicate there
        is no extra information.
        """
        # TODO: Allow the shell argument to be False, True (which uses bash
        #       on Linux and cygwin bash on Windows), command.cmd,
        #       command.cygwin, or command.msys.
        # TODO: Properly determine the cygwin directory on Windows.
        executable = (
            ("bash" if platform.system() != "Windows"
                else "C:/cygwin/cygwin.bat")
            if self.use_shell else None)
        command_arg = (
            self.command if self.use_shell else shlex.split(self.command))

        saved_environ = self.push_environ()

        try:
            return self._run_command_and_get_return_code(
                command_arg, executable), None
        except:
            traceback.print_exc()
            return 1, None
        finally:
            self.pop_environ(saved_environ)

    def _run_command_and_get_return_code(self, command_arg, executable):
        pipe_argument = (subprocess.PIPE if (
            self.stdout_filter is not None or self.stderr_filter is not None)
            else None)
        pipe = subprocess.Popen(
            command_arg, shell=self.use_shell, executable=executable,
            stdout=pipe_argument, stderr=pipe_argument)

        if pipe_argument is None:
            pipe.communicate()
            return pipe.returncode

        pipe_file_to_real_file = {
            pipe.stdout: sys.stdout,
            pipe.stderr: sys.stderr,
        }
        pipe_file_to_filter = {
            pipe.stdout: self.stdout_filter,
            pipe.stderr: self.stderr_filter
        }

        # Print the output as soon as it becomes available, don't merge stdout
        # and stderr together, and preserve the order of the output as much
        # as possible.

        while True:
            if pipe.poll() is not None:
                # The program has exited, so read the rest of the output.
                for pipe_file in (pipe.stdout, pipe.stderr):
                    self._filter_output(
                        pipe_file,
                        pipe_file_to_real_file[pipe_file],
                        pipe_file_to_filter[pipe_file],
                        read_remaining=True)
                break

            # Wait for the program to write some output, and then read in a
            # line of it.  Note that _wait_for_pipe_output can return a file
            # even if a newline has not been written to it, but we'll then call
            # readline on it and block until a newline is written or the end of
            # the file is reached.  If that happens it's ok -- at worst, we'll
            # just delay stdout/stderr output that could have been written
            # earlier.
            readable_pipe_files = self._wait_for_pipe_output(pipe)
            for pipe_file in readable_pipe_files:
                self._filter_output(
                    pipe_file,
                    pipe_file_to_real_file[pipe_file],
                    pipe_file_to_filter[pipe_file])

        # Call the filters with an empty string to indicate that they've
        # reached the end of the output.
        for output_filter in pipe_file_to_filter.values():
            if output_filter is not None:
                output_filter("")

        return pipe.returncode

    def _filter_output(
            self, pipe_file, real_file, output_filter, read_remaining=False):
        """Read a line from the pipe and pass it through the given filter.

        The filter, if not None, is passed a line at a time (including the
        newline character) and returns the filtered output to display.  If
        read_remaining is True, all remaining lines in the pipe are read and
        passed to the filter one at a time.
        """
        while True:
            output = pipe_file.readline()
            if len(output) == 0:
                return

            if output_filter is not None:
                output = output_filter(output)

            if len(output) != 0:
                real_file.write(output)
                real_file.flush()

            if not read_remaining:
                break

    def _wait_for_pipe_output(self, pipe):
        """Given a subprocess.Pipe object, wait for data written to either of
        its stdout or stderr file objects.  Return a tuple of the file objects
        containing data that is waiting to be read.
        """
        if platform.system() != "Windows":
            return select.select(
                (pipe.stdout, pipe.stderr), (), ())[0]

        pipe_files_with_data = []
        for pipe_file in (pipe.stdout, pipe.stderr):
            pipe_handle = msvcrt.get_osfhandle(pipe_file.fileno())
            bytes_available = ctypes.wintypes.DWORD()
            peek_result = win_PeekNamedPipe(
                pipe_handle, None, 0, None, ctypes.byref(bytes_available), None)
            if peek_result == 0:
                continue

            if bytes_available > 0:
                pipe_files_with_data.append(pipe_file)

        return pipe_files_with_data

    def push_environ(self):
        # Remember any environment variables that we'll override.
        saved_environ = dict((name, os.environ[name])
            for name in self.environ if name in os.environ)

        os.environ.update(self.environ)
        return saved_environ

    def pop_environ(self, saved_environ):
        for name in self.environ:
            new_value = saved_environ.get(name)
            if new_value is not None:
                os.environ[name] = new_value
            else:
                del os.environ[name]

def form_command_line(parts):
    return " ".join(_escape_argument(part) for part in parts)

_invalid_char = re.compile(r"[^a-zA-Z0-9\-@%_+=:,./~]")
def _escape_argument(argument):
    if len(argument) == 0:
        return "''"

    if _invalid_char.search(argument) is None:
        return argument

    # Use single quotes and put single quotes into double quotes.
    # For example, the string $'b is then quoted as '$'"'"'b'
    return "'" + argument.replace("'", "'\"'\"'") + "'"

